package dslab.util;

import javax.crypto.*;
import javax.crypto.spec.IvParameterSpec;
import java.io.IOException;
import java.io.PrintWriter;
import java.io.Writer;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;

/**
 * extension of Writer that supports AES communication and can be initialized in a secure
 * session dependent context midway
 * simply wraps writer if encryption is not yet initialized
 */
public class EncryptedPrinter extends Writer {
    private final PrintWriter writer;
    private final Cipher cipher;
    private boolean initialized = false;

    /**
     * EncryptedPrinter constructor
     * @param printer PrintWriter object to wrap
     * @throws NoSuchPaddingException if padding option is invalid
     * @throws NoSuchAlgorithmException if algorithm cannot be used
     */
    public EncryptedPrinter(PrintWriter printer) throws NoSuchPaddingException, NoSuchAlgorithmException {
        this.writer = printer;
        this.cipher = Cipher.getInstance("AES/CTR/NoPadding");
    }

    @Override
    public void write(char[] cbuf, int off, int len) throws IOException {
        writer.write(cbuf, off, len);
    }

    /**
     * takes a string and writes it, simply with `write(String)` of PrintWriter if encryption is not yet
     * initialized and in an encrypted fashion if it is
     * if string is newline terminated (detected using `String.format("%n")`), newline will be removed,
     * text will be encrypted and newline will be re-appended for transmission
     * @param string string to be written
     * @throws IOException if a general I/O failure has occurred
     */
    @Override
    public void write(String string) throws IOException{
        if (initialized) {
            try {
                // if last character is a newline
                if (string.endsWith(String.format("%n")))
                    // strip last character
                    string = String.format("%s%n", CryptoUtil.encryptMessage(string.stripTrailing(), cipher));
                else
                    string = CryptoUtil.encryptMessage(string, cipher);
            } catch (IllegalBlockSizeException | BadPaddingException exception) {
                throw new IOException(exception);
            }
        }
        writer.write(string);
    }

    /**
     * wraps the flush method of the inherent PrintWriter object
     */
    @Override
    public void flush() {
        writer.flush();
    }

    /**
     * wraps the close method of the inherent PrintWriter object
     * @throws IOException if a general I/O failure has occurred
     */
    @Override
    public void close() throws IOException {
        writer.close();
    }

    /**
     * initializes this EncryptedPrinter to switch over to a mode where communication encrypted via AES is used
     * @param key SecretKey to initialize AES communication with
     * @param iv initialization vector to use
     * @throws InvalidAlgorithmParameterException if the encryption algorithm cannot be used
     * @throws InvalidKeyException if the key is not valid
     */
    public void init(SecretKey key, byte[] iv) throws InvalidAlgorithmParameterException, InvalidKeyException {
        cipher.init(Cipher.ENCRYPT_MODE, key, new IvParameterSpec(iv));
        initialized = true;
    }
}
